#batch,使用训练集中的一小部分样本对模型权重进行一次反向传播的参数更新，这一小部分样本被称为“一批数据”
#batch_size，在训练集中选择一小部分来更新权值

import os
import math
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

tf.compat.v1.disable_eager_execution()

#生成图片的路径和标签的list
AJ1 = []
label_AJ1 = []
AJ4 = []
label_AJ4 = []
AJ11 = []
label_AJ11 = []
AJ12 = []
label_AJ12 = []

#首先，对output下的图片路径进行获取，然后存放到对应的列表中，同时对图片再贴上标签
def get_files(file_dir, ratio):
    for file in os.listdir(file_dir + '/AJ1'):
        AJ1.append(file_dir + '/AJ1' + '/' + file)
        label_AJ1.append(0)
    for file in os.listdir(file_dir + '/AJ4'):
        AJ4.append(file_dir + '/AJ4' + '/' + file)
        label_AJ4.append(1)
    for file in os.listdir(file_dir + '/AJ11'):
        AJ11.append(file_dir + '/AJ11' + '/' + file)
        label_AJ11.append(2)
    for file in os.listdir(file_dir + '/AJ12'):
        AJ12.append(file_dir + '/AJ12' + '/' + file)
        label_AJ12.append(3)


    #然后将生成的图片和路径标签list做打乱处理把图像和标签组成一个list
    image_list = np.hstack((AJ1, AJ4, AJ11, AJ12))
    label_list = np.hstack((label_AJ1, label_AJ4, label_AJ11, label_AJ12))

    #利用shuffle打乱顺序
    temp = np.array([image_list, label_list])
    temp = temp.transpose()
    np.random.shuffle(temp)

    #将所有的img和lab转换成list
    all_image_list = list(temp[:,0])
    all_label_list = list(temp[:,1])

    #将所得list分为两部分，一部分用来训练tra，一部分用来测试val
    #ratio是测试集的比例
    n_sample = len(all_label_list)
    n_val = int(math.ceil(n_sample * ratio))  # 测试样本数
    n_train = n_sample - n_val  # 训练样本数

    tra_images = all_image_list[0:n_train]
    tra_labels = all_label_list[0:n_train]
    tra_labels = [int(float(i)) for i in tra_labels]
    val_images = all_image_list[n_train:-1]
    val_labels = all_label_list[n_train:-1]
    val_labels = [int(float(i)) for i in val_labels]

    return tra_images, tra_labels, val_images, val_labels
    
#接下来的部分是用来生成batch
#首先，将上面生成的List传入get_batch() ，转换类型，产生一个输入队列queue，因为img和lab
# 是分开的，所以使用tf.train.slice_input_producer()，然后用tf.read_file()从队列中读取图像
#   image_W, image_H, ：设置好固定的图像高度和宽度
#   设置batch_size：每个batch要放多少张图片
#   capacity：一个队列最大多少
def get_batch(image, label, image_W, image_H, batch_size, capacity):
    # 转换类型
    image = tf.cast(image, tf.string)
    label = tf.cast(label, tf.int32)

    #构建一个输入队列
    input_queue = tf.compat.v1.train.slice_input_producer([image, label])

    label = input_queue[1]
    image_contents = tf.compat.v1.read_file(input_queue[0])  # read img from a queue

    #然后，将图像解码，不同类型的图像不能混在一起，要么只用jpeg，要么只用png等。
    image = tf.image.decode_jpeg(image_contents, channels=3)

    #数据预处理，对图像进行旋转、缩放、裁剪、归一化等操作，让计算出的模型更健壮。
    image = tf.compat.v1.image.resize_image_with_crop_or_pad(image, image_W, image_H)
    image = tf.image.per_image_standardization(image)

    # 最后就是生成batch
    # image_batch: 4D tensor [batch_size, width, height, 3],dtype=tf.float32
    # label_batch: 1D tensor [batch_size], dtype=tf.int32
    image_batch, label_batch = tf.compat.v1.train.batch([image, label],
                                              batch_size=batch_size,
                                              num_threads=32,
                                              capacity=capacity)
    # 重新排列label，行数为[batch_size]
    label_batch = tf.reshape(label_batch, [batch_size])
    #label_batch = tf.convert_to_tensor(label_batch,dtype = tf.float32)
    #label_batch = tf.cast(label_batch,tf.float32)
    image_batch = tf.cast(image_batch, tf.float32)
    return image_batch, label_batch